# -*- encoding: utf-8 -*-
"""
@File: db_handler.py
@Time: 2023/12/27 10:27:07
@Author: cyl
"""
import pymysql
from utils.logger_handle import logger


class DBHandler:
    def __init__(self, host, port, user, password, database, charset, **kwargs):
        self.conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            cursorclass=pymysql.cursors.DictCursor,
            charset=charset,
            **kwargs
        )
        self.cursor = self.conn.cursor()

    def query(self, sql: str, one: bool = True) -> list | dict:
        """
        Returns: [{}, {}, {}...]
        """
        self.cursor.execute(sql)
        try:
            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            logger.error("mysql connect failed", e)
        if one:
            values: tuple = self.cursor.fetchone()
        else:
            values: tuple = self.cursor.fetchall()
        fields: tuple = self.cursor.description
        column_lst: list = []
        for i in fields:
            column_lst.append(i[0])
        sql_lst: list = []
        for row in values:
            result: dict = {}
            for i in range(len(row)):
                result[column_lst[i]] = row[i]
            sql_lst.append(result)
        self.cursor.close()
        self.conn.close()
        
        return sql_lst[0] if len(sql_lst) == 1 else sql_lst
